#!/usr/bin/python2.7
__AUTHOR__ = 'hasherezade'

import argparse
import string
import base64
import ctypes
import random
import hashlib
from Crypto.Cipher import AES

trickbot_b64_charset  = "HJIA/CB+FGKLNOP3RSlUVWXYZfbcdeaghi5kmn0pqrstuvwx89o12467MEDyzQjT"
VERBOSE = False

BS = 16
pad = lambda s: s + (BS - len(s) % BS) * chr(BS - len(s) % BS) 
unpad = lambda s : s[:-ord(s[len(s)-1:])]

def hash_rounds(data_buf):
    while len(data_buf) <= 0x1000:
        buf_hash = hashlib.sha256(data_buf).digest()
        data_buf += buf_hash
    return buf_hash

def aes_decrypt(data):
    key = hash_rounds(data[:0x20])[:0x20]
    iv = hash_rounds(data[0x10:0x30])[:0x10]
    aes = AES.new(key, AES.MODE_CBC, iv)
    data = pad(data[0x30:])
    return aes.decrypt(data)

def dexor(data, key):
    maxlen = len(data)
    keylen = len(key)
    j = 0 #key index
    decoded = ""
    for i in range(0, maxlen):
        kval = key[j % keylen]
        decoded += chr(ord(data[i]) ^ ord(kval))
        j += 1
    return decoded

def to_uint(c):
    return ctypes.c_uint(c).value

def to_int(c):
    return ctypes.c_int(c).value

def str_checksum(b64str):
    v4 = 0
    i = 0
    for c in b64str:
        if c == '=':
            break
        #print "c: %c" % c
        v4 = to_uint(1025 * to_uint((ord(c) + i)))
        #print "v4: %d" % v4
        i = to_int((v4 >> 6) ^ v4)
        #print "i: %d" % i
        #print "---\n"
    step1 = to_uint(9*i) 
    step2 = to_int(9*i ^ (step1 >> 11))
    return to_uint(0x8001 * step2)

def generate_teststring():
    line = ""
    for i in xrange(0, 256):
        line += chr(i)
    return line

def is_charset(line, charset):
    line = line.strip()
    for x in line:
        if x not in charset:
            return False
    return True

def is_printable(line):
    line = line.strip()
    for x in line:
        if x not in string.printable:
            return False
    return True

def trick_base64_dec(s, my_charset):
    if s is None:
        return None
    std_base64chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"
    s = s.translate(string.maketrans(my_charset, std_base64chars))
    data = base64.b64decode(s)
    return data

def trick_base64_enc(s, my_charset):
    if s is None:
        return None
    std_base64chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"
    data = base64.b64encode(s)
    data = data.translate(string.maketrans(std_base64chars, my_charset))
    return data

def swap(c, i, j):
    c = list(c)
    c[i], c[j] = c[j], c[i]
    return ''.join(c)

def randomize_charset(in_charset, n):
    clen = len(in_charset)
    limit = clen-n-1
    while True:
        index1 = random.randint(1,n) + limit
        index2 = random.randint(1,n) + limit
        if index1 != index2:
            break
    in_charset = swap(in_charset, index1, index2)
    return in_charset

def append_padding(val):
    mod = len(val) % 3
    while (len(val) % 4) != 0:
        val += "="
    return val

def process(val, charset):
    decoded = None
    if val is not None and is_charset(val, charset):
        val = append_padding(val)
        try:
            decoded = trick_base64_dec(val, charset)
        except:
            decoded = None
    return decoded

def process_settings(in_lines, charset):
    out_lines = []
    for line in in_lines:
        decoded = process(line, charset)
        if decoded and is_printable(decoded) and len(decoded)>10:
            out_lines.append(decoded)
    return out_lines

def read_settings(fname):
    in_lines = []
    with open(fname) as f:
        for line in f:
            in_lines.append(line)
    return in_lines

def get_botkey(line):
    elements = line.split(' ')
    if elements < 2:
        return None
    return elements[0]

def get_checksum(line):
    elements = line.split(' ')
    if elements < 2:
        return None
    return elements[1]

def make_charset_checksum(my_charset):
    test = generate_teststring()
    test_b64 = trick_base64_enc(test, my_charset)
    test_checks = str_checksum(test_b64)
    return test_checks

def is_valid_checksum(line, test_checks):
    line_checks = get_checksum(line)
    if line_checks is None:
        return False
    if VERBOSE:
        print test_checks , " : ", line_checks
    if line_checks == str(test_checks):
        return True
    return False

def check_charset(in_lines, in_charset):
    test1_checks = make_charset_checksum(in_charset)
    out_lines = process_settings(in_lines, in_charset)
    for line in out_lines:
        #print line
        if is_valid_checksum(line, test1_checks):
            return True
    return False

def remove_spaces(in_lines):
    out_lines2 = []
    for line in in_lines:
        line2 = ""
        for c in line:
            if c == ' ':
                continue
            line2 = line2 + c
        out_lines2.append(line2)
    return out_lines2

def filter_lines(in_lines):
    out_lines = []
    for line in in_lines:
        pair = line.split('=')
        if len(pair) < 2:
            continue
        val  = pair[1].strip()
        for c in val:
            if c == '/':
                out_lines.append(val)
    return remove_spaces(out_lines)


def remove_duplicates(in_lines):
    unique_lines = []
    for line in in_lines:
        if line in unique_lines:
            continue
        unique_lines.append(line)
    return unique_lines
       
def strip_mcconfig(data):
    BEGIN_PATTERN = '<mcconf>'
    END_PATTERN = '</mcconf>'
    begin = data.find(BEGIN_PATTERN)  
    end = data.find(END_PATTERN) + len(END_PATTERN)
    return data[begin:end]

def decode_mcconfig(in_lines, charset, botkey):
    index = 0
    for line in in_lines:
        index += 1
        data = process(line, charset)
        if data is None:
            continue
        data = dexor(data, botkey)
        output = aes_decrypt(data)
        if "mcconf" in output:
            print "\n" + strip_mcconfig(output) + "\n"

def main():
    parser = argparse.ArgumentParser(description="TrickBot settings decoder: decodes botkey from the settings file")
    parser.add_argument('--file', dest="file", default=None, help="TrickBot settings file (i.e. settings.ini)", required=True)
    parser.add_argument('--charset', dest="charset", default=trickbot_b64_charset, help="Charset for Base64 decoding (default: \"" + trickbot_b64_charset + "\")", required=False)
    parser.add_argument('--brute', dest="brute", default=False, help="Brutforce the charset", required=False, action="store_true")
    args = parser.parse_args()

    in_lines = read_settings(args.file)
    in_lines = filter_lines(in_lines)
    in_lines = remove_duplicates(in_lines)

    my_charset = args.charset
    if args.brute:
        print "Searching the charset..."
        while True:
            if check_charset(in_lines, my_charset):
                break
            else:
                my_charset = randomize_charset(my_charset, 9)
                if VERBOSE:
                    print my_charset

    test1_checks = make_charset_checksum(my_charset)
    out_lines = process_settings(in_lines, my_charset)
    for out_line in out_lines:
        if is_valid_checksum(out_line, test1_checks):
            print "[+] Decoded with matching charset: " + my_charset
            decode_mcconfig(in_lines, my_charset, get_botkey(out_line))
        else:
            print "[!] Decoding with basic charset, some characters may be invalid! (Try option --brute)"
        print out_line


if __name__ == "__main__":
    main()

